package controller

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/x509"
	"encoding/base64"
	"encoding/pem"
	"errors"
	"fmt"
	"github.com/gin-gonic/gin"
	"io"
)

// EncryptAES encrypts a plaintext using AES-GCM
func EncryptAES(plaintext, key []byte) (string, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return "", err
	}

	nonce := make([]byte, aes.BlockSize)
	if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
		return "", err
	}

	aesGCM, err := cipher.NewGCM(block)
	if err != nil {
		return "", err
	}

	ciphertext := aesGCM.Seal(nonce, nonce, plaintext, nil)
	return base64.StdEncoding.EncodeToString(ciphertext), nil
}

// EncryptRSALong encrypts a plaintext using RSA-OAEP
func EncryptRSALong(plaintext, publicKeyPEM string) (string, error) {
	block, _ := pem.Decode([]byte(publicKeyPEM))
	if block == nil || block.Type != "RSA PUBLIC KEY" {
		return "", errors.New("failed to decode PEM block containing public key")
	}

	pub, err := x509.ParsePKIXPublicKey(block.Bytes)
	if err != nil {
		return "", err
	}

	rsaPub, ok := pub.(*rsa.PublicKey)
	if !ok {
		return "", errors.New("not a valid RSA public key")
	}

	label := []byte("") // can be empty or a distinctive value
	hash := sha256.New()
	ciphertext, err := rsa.EncryptOAEP(hash, rand.Reader, rsaPub, []byte(plaintext), label)
	if err != nil {
		return "", err
	}

	return base64.StdEncoding.EncodeToString(ciphertext), nil
}

// DecryptRSA decrypts a ciphertext using RSA-OAEP
func DecryptRSA(ciphertextStr, privateKeyPEM string) (string, error) {
	block, _ := pem.Decode([]byte(privateKeyPEM))
	if block == nil || block.Type != "RSA PRIVATE KEY" {
		return "", errors.New("failed to decode PEM block containing private key")
	}

	priv, err := x509.ParsePKCS1PrivateKey(block.Bytes)
	if err != nil {
		return "", err
	}

	ciphertext, err := base64.StdEncoding.DecodeString(ciphertextStr)
	if err != nil {
		return "", err
	}

	label := []byte("") // must match the label used during encryption
	hash := sha256.New()
	plaintext, err := rsa.DecryptOAEP(hash, rand.Reader, priv, ciphertext, label)
	if err != nil {
		return "", err
	}

	return string(plaintext), nil
}

func T(c *gin.Context) {
	publicKeyPEM := `-----BEGIN PUBLIC KEY-----  
MIGfMA0GCSqGSIb3DQEBAQUAA4GNADCBiQKBgQCt7bpQlP7sI0pqgrilzKgZx09s
VzzNXsEp0JOW1QUre3ifMt4R2SWdpYw9Lu0zoJUAIcO3jG43gfKzG2gP3HjmaFib
sgPjw2gxmSV0zY0MAN26gP+sRC5AEnlEz+goSS3JLz2TdMsTYiG0wo053O/a4IdI
nFDuE7JYnEx2po18QwIDAQAB
-----END PUBLIC KEY-----`

	// 哈哈计时（PEM格式）
	privateKeyPEM := `-----BEGIN RSA PRIVATE KEY-----  
MIICeAIBADANBgkqhkiG9w0BAQEFAASCAmIwggJeAgEAAoGBAK3tulCU/uwjSmqC
uKXMqBnHT2xXPM1ewSnQk5bVBSt7eJ8y3hHZJZ2ljD0u7TOglQAhw7eMbjeB8rMb
aA/ceOZoWJuyA+PDaDGZJXTNjQwA3bqA/6xELkASeUTP6ChJLckvPZN0yxNiIbTC
jTnc79rgh0icUO4TslicTHamjXxDAgMBAAECgYALUPRWWdtcMc+drnOlFfO5+aXJ
PsPN+ORIfDGeFq4Hbjki4oRgVOiFoUfGk08LlN8c6xV408I4ClRoIthab9XNzHA4
dvGMANqD9FS+EAjvWj1NbqB7i8NcTbzUET3I/JsDII4JocrP4HTsVp0XBm2VxqM9
81Skbn1fKkMNHh2QgQJBAMoBvZWovYHispa2Bk8UBngqimHREkIXhd8C+7Tn+SV5
g+BCDbKO1TKVlbygSihD7gZwYQVJXoTuZ/4zXrQGwUkCQQDcar7xAh3yXGHHy/xg
RiWGBNvVV0vx0wBCSIx2oHh7W+lbZn+GaUp6L52cLHpNIViqeogYjg8hJg/b7HZ/
F90rAkEAmuvMzrsTY0pNtmGw4vKzp3LQoAXsX72UkJzEF1CUddYEeneN87Ow89Ne
e8PeO5UlNdD3x0smbwI7xOr2pVBM6QJBAK/NtF4mghw88wCZAsMHMAak+XJAHqLS
2PIFS4OSmoZJ2Oa96hkbAfxS3jMihqneXwoakVKEfCzEXwHaqOjV7+ECQQCo2IbX
KUj6m67+YdXCZc9JkDeck91Y5cd/Ze2Wx7LrvhBXLLm9PKLpmWOOVGAUy9APZtwE
IwQTyCq/Ylmq+I+m
-----END RSA PRIVATE KEY-----`
	// Long string to encrypt
	longString := "This is a very long string that we want to encrypt using AES and then encrypt the AES key using RSA."

	// Generate AES key
	aesKey := make([]byte, 32) // AES-256
	if _, err := io.ReadFull(rand.Reader, aesKey); err != nil {
		fmt.Println("Error generating AES key:", err)
		return
	}

	// Encrypt long string using AES
	aesCiphertext, err := EncryptAES([]byte(longString), aesKey)
	if err != nil {
		fmt.Println("Error encrypting long string with AES:", err)
		return
	}
	fmt.Println("AES Ciphertext:", aesCiphertext)

	// Encrypt AES key using RSA
	rsaCiphertext, err := EncryptRSALong(string(aesKey), publicKeyPEM)
	if err != nil {
		fmt.Println("Error encrypting AES key with RSA:", err)
		return
	}
	fmt.Println("RSA Ciphertext (AES Key):", rsaCiphertext)

	// Decrypt AES key using RSA
	decryptedAESKey, err := DecryptRSA(rsaCiphertext, privateKeyPEM)
	if err != nil {
		fmt.Println("Error decrypting AES key with RSA:", err)
		return
	}
	fmt.Println("Decrypted AES Key:", decryptedAESKey)
}
